# -*- coding: utf-8 -*-
# @Author: Tim Liu
# @Date: 2025-01-09
# @Last Modified by: Cursor
# @Last Modified time: 2025-01-13

# @Description: Schema Graph Transformer

import json
from typing import Dict, List, Tuple
from langchain_experimental.graph_transformers import LLMGraphTransformer
from langchain_core.prompts import ChatPromptTemplate

system_prompt = (
    "# Knowledge Graph Instructions for GPT-4\n"
    "## 1. Overview\n"
    "You are a top-tier algorithm designed for extracting information in structured "
    "formats to build a knowledge graph.\n"
    "Try to capture as much information from the text as possible without "
    "sacrificing accuracy. Do not add any information that is not explicitly "
    "mentioned in the text\n"
    "- **Nodes** represent entities and concepts.\n"
    "- The aim is to achieve simplicity and clarity in the knowledge graph, making it\n"
    "accessible for a vast audience.\n"
    "## 2. Labeling Nodes\n"
    "- **Consistency**: Ensure you use available types for node labels.\n"
    "Ensure you use basic or elementary types for node labels.\n"
    "- For example, when you identify an entity representing a person, "
    "always label it as **'person'**. Avoid using more specific terms "
    "like 'mathematician' or 'scientist'.\n"
    "  - **Node IDs**: Node IDs should be unique and contain"
    "names or human-readable identifiers found in the text. Node IDs shall include hierarchical "
    "identifiers for sub-objects. For example, if an object of Type has relationship HAS_SubType or relationship with prefix HAS_ , "
    "then use compound Node IDs for SubType. For example, for relationship Type -HAS_SubType- SubType, the Node ID of SubType shall be Type_SubType. "
    "Please do NOT generate same Node ID for several sub-objects."
    "- **Relationships** represent connections between entities or concepts.\n"
    "Ensure consistency and generality in relationship types when constructing "
    "knowledge graphs. Instead of using specific and momentary types "
    "such as 'BECAME_PROFESSOR', use more general and timeless relationship types "
    "like 'PROFESSOR'. Make sure to use general and timeless relationship types!\n"
    "## 3. Coreference Resolution\n"
    "- **Maintain Entity Consistency**: When extracting entities, it's vital to "
    "ensure consistency.\n"
    'If an entity, such as "John Doe", is mentioned multiple times in the text '
    'but is referred to by different names or pronouns (e.g., "Joe", "he"),'
    "always use the most complete identifier for that entity throughout the "
    'knowledge graph. In this example, use "John Doe" as the entity ID.\n'
    "Remember, the knowledge graph should be coherent and easily understandable, "
    "so maintaining consistency in entity references is crucial.\n"
    "## 4. Strict Compliance\n"
    "Adhere to the rules strictly. Non-compliance will result in termination."
    "## 5. source_url and page\n"
    "Try your best to extract source_url and page properties for node, these factual information are important."
)

custom_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            system_prompt,
        ),
        (
            "human",
            (
                "Tip: Make sure to answer in the correct format and do "
                "not include any explanations. Do NOT make up graph nodes or relationships if not mentioned in below input. "
                "If the input is too short to know what is the substance, please reply I do not know"
                "Use the given format to extract information from the "
                "following input: {input}"
            ),
        ),
    ]
)

class SchemaGraphTransformer(LLMGraphTransformer):
    def __init__(self, llm, allowed_nodes, allowed_relationships, node_properties, prompt=custom_prompt):
        super().__init__(llm=llm, allowed_nodes=allowed_nodes, allowed_relationships=allowed_relationships, node_properties=node_properties, prompt=prompt)
    
    @staticmethod
    def convert_schema_to_graphmeta(schema: Dict) -> Tuple[List[str], List[Tuple[str, str, str]], List[str], Dict[str, Dict[str, List[str]]]]:
        """
        Converts a schema JSON to LLMGraphTransformer compatible allowed_nodes, allowed_relationships, node_properties, and indexes.

        Args:
            schema (Dict): The schema JSON.

        Returns:
            Tuple[List[str], List[Tuple[str, str, str]], List[str], Dict[str, Dict[str, List[str]]]]: allowed_nodes, allowed_relationships, node_properties, indexes.
        """
        allowed_nodes = list(schema["node_types"].keys())
        allowed_relationships = []
        node_properties = set()
        indexes = {}

        for rel_type, rel_info in schema["relationships"].items():
            source = rel_info["source"]
            target = rel_info["target"]
            allowed_relationships.append((source, rel_type, target))

        for node_type, node_info in schema["node_types"].items():
            for prop in node_info["properties"].keys():
                node_properties.add(prop)
            if "indexes" in node_info:
                indexes[node_type] = node_info["indexes"]            

        return allowed_nodes, allowed_relationships, list(node_properties), indexes

    @staticmethod
    def convert_graphmeta_to_schema(allowed_nodes: List[str], allowed_relationships: List[Tuple[str, str, str]], node_properties: List[str], indexes: Dict[str, Dict[str, Dict[str, str]]]) -> str:
        """
        Converts LLMGraphTransformer compatible allowed_nodes, allowed_relationships, node_properties, and indexes back to a schema JSON string.

        Args:
            allowed_nodes (List[str]): List of allowed node types.
            allowed_relationships (List[Tuple[str, str, str]]): List of allowed relationships.
            node_properties (List[str]): List of node properties.
            indexes (Dict[str, Dict[str, Dict[str, str]]]): Dictionary of indexes.

        Returns:
            str: The schema JSON string.
        """
        schema = {
            "node_types": {node: {"properties": {prop: "STRING" for prop in node_properties}} for node in allowed_nodes},
            "relationships": {rel[1]: {"source": rel[0], "target": rel[2], "properties": {}} for rel in allowed_relationships}
        }

        for node_type, index_info in indexes.items():
            if node_type in schema["node_types"]:
                schema["node_types"][node_type]["indexes"] = index_info
                
        return json.dumps(schema)